from action import *
from state import *
from goalparser import *
from plannersupport import is_discrete
from parserbasic import tokenize_code

class PlanningProblem:
	def __init__(self):
		self.env = dict()
		self.actions = []
		self.actionSamplers = []

	def updateEnv(self,env):
		self.env.update(env);

	def setInitial(self,stateSpace,initial):
		self.stateSpace = stateSpace
		self.initialState = initial

	def addAction(self,action):
		self.actions.append(action)
		self.actionSamplers.append(ActionParameterSampler(action,self.env))

	def setGoalTest(self,goal):
		self.goalTest = goal
	
	def setStaticFeatures(self, features):
		self.staticFeatures = features
		for f in features:
			self.initialState = features[f].ex(self.env, self.initialState)
		
	def setDynamicFeatures(self, features):
		self.dynamicFeatures = features
		for f in features:
			self.initialState = features[f].ex(self.env, self.initialState)

	def print_properties(self):
		print "%d state objects, %d state constraints, %d actions, %d goal conditions"%(len(self.stateSpace.names),len(self.stateSpace.constraints), len(self.actions),len(self.goalTest.tests))
		print "State variables:",', '.join(self.stateSpace.names)
		for a in self.actions:
			print '  Action',a.name
			print '	%d parameters, %d preconditions, %d effects'%(len(a.params),len(a.preconds),len(a.effects))		

	def constructDiscrete(self):
		"""Constructs a version of the planning problem with all continuous
		variables stripped out"""
		pnew = PlanningProblem()
		pnew.env = self.env
		discreteStateVars = []
		continuousStateVars = []
		for (k,v) in self.initialState.iteritems():
			if is_discrete(v):
				discreteStateVars.append(k)
			else:
				continuousStateVars.append(k)
		discreteStateVars = set(discreteStateVars)
		continuousStateVars = set(continuousStateVars)

		newstateconstraints = []
		for c in self.stateSpace.constraints:
			deps = set(tokenize_code(c))
			if len(deps & continuousStateVars) == 0:
				newstateconstraints.append(c)
			
		pnew.setInitial(StateSpace(discreteStateVars,newstateconstraints),dict([(k,v) for (k,v) in self.initialState.iteritems() if k in discreteStateVars]))

		print 'Discrete states',list(discreteStateVars)

		for a in self.actions:
			anew = Action(a.name)
			continuousParams = set(a.continuousParams())
			discreteParams = set(a.discreteParams())
			delayParams = set(a.delayParams())
			for p in delayParams:
				allDiscrete = True
				for dep in a.params[p].deps:
					if dep in continuousStateVars or dep in continuousParams:
						allDiscrete=False
						break
				if allDiscrete:
					discreteParams.add(p)
				else:
					continuousParams.add(p)
			print 'Action',a.name,'params',discreteParams
			for p in discreteParams:
				print p,a.params[p].__class__
				anew.addParam(p,a.params[p])
			for pre in a.normalPreconds():
				if (pre.paramdeps.issubset(discreteParams)) and (pre.statedeps.issubset(discreteStateVars)):
					anew.addPrecond(pre)
			for e in a.effects:
				deps = set(tokenize_code(e))
				if len(deps & continuousStateVars) == 0 and len(deps & continuousParams) == 0:
					anew.addEffect(e)
			#eliminate continous-dependent costs
			anew.costTerms = []
			for cost in a.costTerms:
				if isinstance(cost,(int,float)):
					anew.costTerms.append(cost)
				elif isinstance(cost,str):
					deps = set(tokenize_code(cost))
					if len(deps & continuousStateVars) == 0 and len(deps&continuousParams)==0:
						anew.costTerms.append(cost)

			#if not empty add it
			if len(a.effects)>0:
				pnew.addAction(anew)

		newgoals = []
		for g in self.goalTest.tests:
			deps = set(tokenize_code(g))
			if len(deps & continuousStateVars) == 0:
				newgoals.append(g)
		pnew.setGoalTest(GoalTest(newgoals))
		return pnew

	def bind_globals(self):
		"""Binds import files to be used in PyTAMP files.
		Loads up domain.py, cspace.py, and iterative_subproblem.py.
		"""
		self.env = {}
		#Place all the default code available to the user here
		exec 'from domain import *' in self.env
		exec 'from cspace import *' in self.env
		exec 'from iterative_subproblem import *' in self.env

	def bind_builtins(self):
		"""Binds some default functions into the PyTAMP namespace

		- ACTION_NAME(u1=val1,...,un=valn) produces a bound action
		  from keyword arguments. (No keyword checking is performed!)
		  
		- apply(action,x) applies the bound action to x
		
		- applicable(action,x) returns true if the action is applicable
		
		- resolve(partialaction,x,iters) resolves the unbound action
		  parameters on the state x using at most iters iterations.
		  Currently, partial actions may not include vector params.
		  The bound action is returned.
		  """

		class bind:
			def __init__(self,action):
				self.action = action
			def __call__(self,**kwparams):
				return (self.action,kwparams)
		class apply:
			def __init__(self,env):
				self.env = env
			def __call__(self,action,state):
				return action[0].apply(self.env,action[1],state)
		class applicable:
			def __init__(self,env):
				self.env = env
			def __call__(self,action,state):
				return action[0].applicable(self.env,action[1],state)
		class cost:
			def __init__(self,env):
				self.env = env
			def __call__(self,action,state=None):
				return action[0].cost(self.env,action[1],state)
		class resolve:
			def __init__(self,problem):
				self.problem = problem
			def __call__(self,partialaction,x,iters):
				sampler = None
				for s in self.problem.actionSamplers:
					if s.action == partialaction[0]:
						sampler = s
						break
				if sampler == None:
					raise RuntimeError("Unable to find action"+partialaction[0].name)
				sampler.bindState(x)
				newparams = sampler.sample_applicable_partial(partialaction[1],iters)
				if newparams == None: return None
				
				params = partialaction[1].copy()
				params.update(newparams)
				return (partialaction[0],params)

		for a in self.actions:
			self.env[a.name] = bind(a)

		self.env['apply'] = apply(self.env)
		self.env['applicable'] = applicable(self.env)
		self.env['cost'] = cost(self.env)
		self.env['resolve'] = resolve(self)
							 
